{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nUse Tensor Expression Debug Display (TEDD) for Visualization\n============================================================\n**Author**: `Yongfeng Gu <https://github.com/yongfeng-nv>`_\n\nThis is an introduction about using TEDD to visualize tensor expressions.\n\nTensor Expressions are scheduled with primitives.  Although individual\nprimitives are usually easy to understand, they become complicated quickly\nwhen you put them together. We have introduced an operational model of\nschedule primitives in Tensor Expression.\n\n* the interactions between different schedule primitives,\n* the impact of the schedule primitives on the final code generation.\n\nThe operational model is based on a Dataflow Graph, a Schedule Tree and an\nIterVar Relationship Graph. Schedule primitives perform operations on these\ngraphs.\n\nTEDD renders these three graphs from a given schedule.  This tutorial demonstrates\nhow to use TEDD and how to interpret the rendered graphs.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tvm\nfrom tvm import te\nfrom tvm import topi\nfrom tvm.contrib import tedd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define and Schedule Convolution with Bias and ReLU\n--------------------------------------------------\nLet's build an example Tensor Expression for a convolution followed by Bias and ReLU.\nWe first connect conv2d, add, and relu TOPIs.  Then, we create a TOPI generic schedule.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "batch = 1\nin_channel = 256\nin_size = 32\nnum_filter = 256\nkernel = 3\nstride = 1\npadding = \"SAME\"\ndilation = 1\n\nA = te.placeholder((in_size, in_size, in_channel, batch), name=\"A\")\nW = te.placeholder((kernel, kernel, in_channel, num_filter), name=\"W\")\nB = te.placeholder((1, num_filter, 1), name=\"bias\")\n\nwith tvm.target.Target(\"llvm\"):\n    t_conv = topi.nn.conv2d_hwcn(A, W, stride, padding, dilation)\n    t_bias = topi.add(t_conv, B)\n    t_relu = topi.nn.relu(t_bias)\n    s = topi.generic.schedule_conv2d_hwcn([t_relu])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Render Graphs with TEDD\n-----------------------\nWe render graphs to see the computation\nand how it is scheduled.\nIf you run the tutorial in a Jupyter notebook, you can use the following commented lines\nto render SVG figures showing in notebook directly.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tedd.viz_dataflow_graph(s, dot_file_path=\"/tmp/dfg.dot\")\n# tedd.viz_dataflow_graph(s, show_svg = True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "![](https://github.com/dmlc/web-data/raw/main/tvm/tutorial/tedd_dfg.png)\n\n     :align: center\n\nThe first one is a dataflow graph.  Every node represents a stage with name and memory\nscope shown in the middle and inputs/outputs information on the sides.\nEdges show nodes' dependency.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tedd.viz_schedule_tree(s, dot_file_path=\"/tmp/scheduletree.dot\")\n# tedd.viz_schedule_tree(s, show_svg = True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We just rendered the schedule tree graph.  You may notice an warning about ranges not\navailable.\nThe message also suggests to call normalize() to infer range information.  We will\nskip inspecting the first schedule tree and encourage you to compare the graphs before\nand after normalize() for its impact.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "s = s.normalize()\ntedd.viz_schedule_tree(s, dot_file_path=\"/tmp/scheduletree2.dot\")\n# tedd.viz_schedule_tree(s, show_svg = True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "![](https://github.com/dmlc/web-data/raw/main/tvm/tutorial/tedd_st.png)\n\n     :align: center\n\nNow, let us take a close look at the second schedule tree.  Every block under ROOT\nrepresents a\nstage.  Stage name shows in the top row and compute shows in the bottom row.\nThe middle rows are for IterVars, the higher the outer, the lower the inner.\nAn IterVar row contains its index, name, type, and other optional information.\nLet's use the W.shared stage as an example.  The top row tells\nits name, \"W.shared\", and memory scope, \"Shared\".  Its compute is\n:code:`W(ax0, ax1, ax2, ax3)`.\nIts outer most loop IterVar is ax0.ax1.fused.ax2.fused.ax3.fused.outer,\nindexed with 0, of kDataPar, bound to threadIdx.y, and with range(min=0, ext=8).\nYou can also tell\nIterVar type with the index box color, shown in the legend.\n\nIf a stage doesn't compute_at any other stage, it has an edge directly to the\nROOT node.  Otherwise, it has an edge pointing to the IterVar it attaches to,\nsuch as W.shared attaches to rx.outer in the middle compute stage.\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>By definition, IterVars are internal nodes and computes are leaf nodes in\n  a schedule tree.   The edges among IterVars and compute within one stage are\n  omitted, making every stage a block, for better readability.</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tedd.viz_itervar_relationship_graph(s, dot_file_path=\"/tmp/itervar.dot\")\n# tedd.viz_itervar_relationship_graph(s, show_svg = True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "![](https://github.com/dmlc/web-data/raw/main/tvm/tutorial/tedd_itervar_rel.png)\n\n     :align: center\n\nThe last one is an IterVar Relationship Graph.  Every subgraph represents a\nstage and contains IterVar nodes and transformation nodes.  For example,\nW.shared has three split nodes and three fuse nodes.  The rest are IterVar\nnodes of the same format as the IterVar rows in Schedule Trees.  Root\nIterVars are those not driven by any transformation node, such as ax0; leaf\nIterVars don't drive any transformation node and have non-negative indices,\nsuch as ax0.ax1.fused.ax2.fused.ax3.fused.outer with index of 0.\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Summary\n-------\nThis tutorial demonstrates the usage of TEDD.  We use an example built\nwith TOPI to show the schedules under the hood.  You can also use\nit before and after any schedule primitive to inspect its effect.\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}